// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <cstring>
#include <memory>
#include <mutex>

#include <nebula/array/array_base.h>
#include <nebula/array/array_dict.h>
#include <nebula/array/array_nested.h>
#include <nebula/array/builder_primitive.h>
#include <nebula/array/concatenate.h>
#include <nebula/array/dict_internal.h>
#include <nebula/array/util.h>
#include <nebula/core/buffer.h>
#include <nebula/compute/api_vector.h>
#include <nebula/compute/cast.h>
#include <nebula/compute/kernels/common_internal.h>

#include <nebula/util/hashing.h>
#include <nebula/numeric/int_util.h>


namespace nebula {

    using internal::DictionaryTraits;
    using internal::HashTraits;
    using internal::TransposeInts;
}
namespace nebula::compute::internal {

    namespace {

        class ActionBase {
        public:
            ActionBase(const std::shared_ptr<DataType> &type, MemoryPool *pool)
                    : type_(type), pool_(pool) {}

        protected:
            std::shared_ptr<DataType> type_;
            MemoryPool *pool_;
        };

// ----------------------------------------------------------------------
// Unique

        class UniqueAction final : public ActionBase {
        public:
            using ActionBase::ActionBase;

            static constexpr bool with_error_status = false;

            UniqueAction(const std::shared_ptr<DataType> &type, const FunctionOptions *options,
                         MemoryPool *pool)
                    : ActionBase(type, pool) {}

            turbo::Status reset() { return turbo::OkStatus(); }

            turbo::Status Reserve(const int64_t length) { return turbo::OkStatus(); }

            template<class Index>
            void ObserveNullFound(Index index) {}

            template<class Index>
            void ObserveNullNotFound(Index index) {}

            template<class Index>
            void ObserveFound(Index index) {}

            template<class Index>
            void ObserveNotFound(Index index) {}

            bool ShouldEncodeNulls() { return true; }

            turbo::Status flush(ExecResult *out) { return turbo::OkStatus(); }

            turbo::Status FlushFinal(ExecResult *out) { return turbo::OkStatus(); }
        };

// ----------------------------------------------------------------------
// Count values

        class ValueCountsAction final : ActionBase {
        public:
            using ActionBase::ActionBase;

            static constexpr bool with_error_status = true;

            ValueCountsAction(const std::shared_ptr<DataType> &type, const FunctionOptions *options,
                              MemoryPool *pool)
                    : ActionBase(type, pool), count_builder_(pool) {}

            turbo::Status Reserve(const int64_t length) {
                // builder size is independent of input array size.
                return turbo::OkStatus();
            }

            turbo::Status reset() {
                count_builder_.reset();
                return turbo::OkStatus();
            }

            // Don't do anything on flush because we don't want to finalize the builder
            // or incur the cost of memory copies.
            turbo::Status flush(ExecResult *out) { return turbo::OkStatus(); }

            // Return the counts corresponding the MemoTable keys.
            turbo::Status FlushFinal(ExecResult *out) {
                std::shared_ptr<ArrayData> result;
                TURBO_RETURN_NOT_OK(count_builder_.finish_internal(&result));
                out->value = std::move(result);
                return turbo::OkStatus();
            }

            template<class Index>
            void ObserveNullFound(Index index) {
                count_builder_[index]++;
            }

            template<class Index>
            void ObserveNullNotFound(Index index) {
                KLOG(FATAL) << "ObserveNullNotFound without err_status should not be called";
            }

            template<class Index>
            void ObserveNullNotFound(Index index, turbo::Status *status) {
                turbo::Status s = count_builder_.append(1);
                if (TURBO_UNLIKELY(!s.ok())) {
                    *status = s;
                }
            }

            template<class Index>
            void ObserveFound(Index slot) {
                count_builder_[slot]++;
            }

            template<class Index>
            void ObserveNotFound(Index slot, turbo::Status *status) {
                turbo::Status s = count_builder_.append(1);
                if (TURBO_UNLIKELY(!s.ok())) {
                    *status = s;
                }
            }

            bool ShouldEncodeNulls() const { return true; }

        private:
            Int64Builder count_builder_;
        };

// ----------------------------------------------------------------------
// Dictionary encode implementation

        class DictEncodeAction final : public ActionBase {
        public:
            using ActionBase::ActionBase;

            static constexpr bool with_error_status = false;

            DictEncodeAction(const std::shared_ptr<DataType> &type, const FunctionOptions *options,
                             MemoryPool *pool)
                    : ActionBase(type, pool), indices_builder_(pool) {
                if (auto options_ptr = static_cast<const DictionaryEncodeOptions *>(options)) {
                    encode_options_ = *options_ptr;
                }
            }

            turbo::Status reset() {
                indices_builder_.reset();
                return turbo::OkStatus();
            }

            turbo::Status Reserve(const int64_t length) { return indices_builder_.Reserve(length); }

            template<class Index>
            void ObserveNullFound(Index index) {
                if (encode_options_.null_encoding_behavior == DictionaryEncodeOptions::MASK) {
                    indices_builder_.unsafe_append_null();
                } else {
                    indices_builder_.unsafe_append(index);
                }
            }

            template<class Index>
            void ObserveNullNotFound(Index index) {
                ObserveNullFound(index);
            }

            template<class Index>
            void ObserveFound(Index index) {
                indices_builder_.unsafe_append(index);
            }

            template<class Index>
            void ObserveNotFound(Index index) {
                ObserveFound(index);
            }

            bool ShouldEncodeNulls() {
                return encode_options_.null_encoding_behavior == DictionaryEncodeOptions::ENCODE;
            }

            turbo::Status flush(ExecResult *out) {
                std::shared_ptr<ArrayData> result;
                TURBO_RETURN_NOT_OK(indices_builder_.finish_internal(&result));
                out->value = std::move(result);
                return turbo::OkStatus();
            }

            turbo::Status FlushFinal(ExecResult *out) { return turbo::OkStatus(); }

        private:
            Int32Builder indices_builder_;
            DictionaryEncodeOptions encode_options_;
        };

        class HashKernel : public KernelState {
        public:
            HashKernel() : options_(nullptr) {}

            explicit HashKernel(const FunctionOptions *options) : options_(options) {}

            // reset for another run.
            virtual turbo::Status reset() = 0;

            // flush out accumulated results from the last invocation of Call.
            virtual turbo::Status flush(ExecResult *out) = 0;

            // flush out accumulated results across all invocations of Call. The kernel
            // should not be used until after reset() is called.
            virtual turbo::Status FlushFinal(ExecResult *out) = 0;

            // Get the values (keys) accumulated in the dictionary so far.
            virtual turbo::Status GetDictionary(std::shared_ptr<ArrayData> *out) = 0;

            virtual std::shared_ptr<DataType> get_value_type() const = 0;

            turbo::Status append(KernelContext *ctx, const ArraySpan &input) {
                std::lock_guard<std::mutex> guard(lock_);
                return append(input);
            }

            // Prepare the Action for the given input (e.g. reserve appropriately sized
            // data structures) and visit the given input with Action.
            virtual turbo::Status append(const ArraySpan &arr) = 0;

        protected:
            const FunctionOptions *options_;
            std::mutex lock_;
        };

// ----------------------------------------------------------------------
// Base class for all "regular" hash kernel implementations
// (NullType has a separate implementation)

        template<typename Type, typename Action, typename Scalar = typename Type::c_type,
                bool with_error_status = Action::with_error_status>
        class RegularHashKernel : public HashKernel {
        public:
            RegularHashKernel(const std::shared_ptr<DataType> &type, const FunctionOptions *options,
                              MemoryPool *pool)
                    : HashKernel(options), pool_(pool), type_(type), action_(type, options, pool) {}

            turbo::Status reset() override {
                memo_table_.reset(new MemoTable(pool_, 0));
                return action_.reset();
            }

            turbo::Status append(const ArraySpan &arr) override {
                TURBO_RETURN_NOT_OK(action_.Reserve(arr.length));
                return DoAppend(arr);
            }

            turbo::Status flush(ExecResult *out) override { return action_.flush(out); }

            turbo::Status FlushFinal(ExecResult *out) override { return action_.FlushFinal(out); }

            turbo::Status GetDictionary(std::shared_ptr<ArrayData> *out) override {
                TURBO_MOVE_OR_RAISE(*out, DictionaryTraits<Type>::GetDictionaryArrayData(
                        pool_, type_, *memo_table_, 0 /* start_offset */));
                return turbo::OkStatus();
            }

            std::shared_ptr<DataType> get_value_type() const override { return type_; }

            template<bool HasError = with_error_status>
            enable_if_t<!HasError, turbo::Status> DoAppend(const ArraySpan &arr) {
                return VisitArraySpanInline<Type>(
                        arr,
                        [this](Scalar v) {
                            auto on_found = [this](int32_t memo_index) {
                                action_.ObserveFound(memo_index);
                            };
                            auto on_not_found = [this](int32_t memo_index) {
                                action_.ObserveNotFound(memo_index);
                            };

                            int32_t unused_memo_index;
                            return memo_table_->GetOrInsert(v, std::move(on_found), std::move(on_not_found),
                                                            &unused_memo_index);
                        },
                        [this]() {
                            if (action_.ShouldEncodeNulls()) {
                                auto on_found = [this](int32_t memo_index) {
                                    action_.ObserveNullFound(memo_index);
                                };
                                auto on_not_found = [this](int32_t memo_index) {
                                    action_.ObserveNullNotFound(memo_index);
                                };
                                memo_table_->GetOrInsertNull(std::move(on_found), std::move(on_not_found));
                            } else {
                                action_.ObserveNullNotFound(-1);
                            }
                            return turbo::OkStatus();
                        });
            }

            template<bool HasError = with_error_status>
            enable_if_t<HasError, turbo::Status> DoAppend(const ArraySpan &arr) {
                return VisitArraySpanInline<Type>(
                        arr,
                        [this](Scalar v) {
                            turbo::Status s = turbo::OkStatus();
                            auto on_found = [this](int32_t memo_index) {
                                action_.ObserveFound(memo_index);
                            };
                            auto on_not_found = [this, &s](int32_t memo_index) {
                                action_.ObserveNotFound(memo_index, &s);
                            };

                            int32_t unused_memo_index;
                            TURBO_RETURN_NOT_OK(memo_table_->GetOrInsert(
                                    v, std::move(on_found), std::move(on_not_found), &unused_memo_index));
                            return s;
                        },
                        [this]() {
                            // Null
                            turbo::Status s = turbo::OkStatus();
                            auto on_found = [this](int32_t memo_index) {
                                action_.ObserveNullFound(memo_index);
                            };
                            auto on_not_found = [this, &s](int32_t memo_index) {
                                action_.ObserveNullNotFound(memo_index, &s);
                            };
                            if (action_.ShouldEncodeNulls()) {
                                memo_table_->GetOrInsertNull(std::move(on_found), std::move(on_not_found));
                            }
                            return s;
                        });
            }

        protected:
            using MemoTable = typename HashTraits<Type>::MemoTableType;

            MemoryPool *pool_;
            std::shared_ptr<DataType> type_;
            Action action_;
            std::unique_ptr<MemoTable> memo_table_;
        };

// ----------------------------------------------------------------------
// Hash kernel implementation for nulls

        template<typename Action, bool with_error_status = Action::with_error_status>
        class NullHashKernel : public HashKernel {
        public:
            NullHashKernel(const std::shared_ptr<DataType> &type, const FunctionOptions *options,
                           MemoryPool *pool)
                    : pool_(pool), type_(type), action_(type, options, pool) {}

            turbo::Status reset() override { return action_.reset(); }

            turbo::Status append(const ArraySpan &arr) override { return DoAppend(arr); }

            template<bool HasError = with_error_status>
            enable_if_t<!HasError, turbo::Status> DoAppend(const ArraySpan &arr) {
                TURBO_RETURN_NOT_OK(action_.Reserve(arr.length));
                for (int64_t i = 0; i < arr.length; ++i) {
                    if (i == 0) {
                        seen_null_ = true;
                        action_.ObserveNullNotFound(0);
                    } else {
                        action_.ObserveNullFound(0);
                    }
                }
                return turbo::OkStatus();
            }

            template<bool HasError = with_error_status>
            enable_if_t<HasError, turbo::Status> DoAppend(const ArraySpan &arr) {
                turbo::Status s = turbo::OkStatus();
                TURBO_RETURN_NOT_OK(action_.Reserve(arr.length));
                for (int64_t i = 0; i < arr.length; ++i) {
                    if (seen_null_ == false && i == 0) {
                        seen_null_ = true;
                        action_.ObserveNullNotFound(0, &s);
                    } else {
                        action_.ObserveNullFound(0);
                    }
                }
                return s;
            }

            turbo::Status flush(ExecResult *out) override { return action_.flush(out); }

            turbo::Status FlushFinal(ExecResult *out) override { return action_.FlushFinal(out); }

            turbo::Status GetDictionary(std::shared_ptr<ArrayData> *out) override {
                std::shared_ptr<NullArray> null_array;
                if (seen_null_) {
                    null_array = std::make_shared<NullArray>(1);
                } else {
                    null_array = std::make_shared<NullArray>(0);
                }
                *out = null_array->data();
                return turbo::OkStatus();
            }

            std::shared_ptr<DataType> get_value_type() const override { return type_; }

        protected:
            MemoryPool *pool_;
            std::shared_ptr<DataType> type_;
            bool seen_null_ = false;
            Action action_;
        };

// ----------------------------------------------------------------------
// Hashing for dictionary type

        class DictionaryHashKernel : public HashKernel {
        public:
            explicit DictionaryHashKernel(std::unique_ptr<HashKernel> indices_kernel,
                                          std::shared_ptr<DataType> dictionary_value_type)
                    : indices_kernel_(std::move(indices_kernel)),
                      dictionary_value_type_(std::move(dictionary_value_type)) {}

            turbo::Status reset() override { return indices_kernel_->reset(); }

            turbo::Status append(const ArraySpan &arr) override {
                auto arr_dict = arr.dictionary().to_array();
                if (!first_dictionary_) {
                    first_dictionary_ = arr_dict;
                } else if (!first_dictionary_->equals(*arr_dict)) {
                    // NOTE: This approach computes a new dictionary unification per chunk.
                    // This is in effect O(n*k) where n is the total chunked array length and
                    // k is the number of chunks (therefore O(n**2) if chunks have a fixed size).
                    //
                    // A better approach may be to run the kernel over each individual chunk,
                    // and then hash-aggregate all results (for example sum-group-by for
                    // the "value_counts" kernel).
                    if (dictionary_unifier_ == nullptr) {
                        TURBO_MOVE_OR_RAISE(dictionary_unifier_,
                                            DictionaryUnifier::create(first_dictionary_->type()));
                        TURBO_RETURN_NOT_OK(dictionary_unifier_->unify(*first_dictionary_));
                    }
                    auto out_dict_type = first_dictionary_->type();
                    std::shared_ptr<Buffer> transpose_map;

                    TURBO_RETURN_NOT_OK(dictionary_unifier_->unify(*arr_dict, &transpose_map));

                    auto transpose = reinterpret_cast<const int32_t *>(transpose_map->data());
                    auto in_array = arr.to_array();
                    const auto &in_dict_array =
                            turbo::checked_cast<const DictionaryArray &>(*in_array);
                    TURBO_MOVE_OR_RAISE(
                            auto tmp, in_dict_array.transpose(arr.type->get_shared_ptr(),
                                                              in_dict_array.dictionary(), transpose));
                    return indices_kernel_->append(*tmp->data());
                }

                return indices_kernel_->append(arr);
            }

            turbo::Status flush(ExecResult *out) override { return indices_kernel_->flush(out); }

            turbo::Status FlushFinal(ExecResult *out) override { return indices_kernel_->FlushFinal(out); }

            turbo::Status GetDictionary(std::shared_ptr<ArrayData> *out) override {
                return indices_kernel_->GetDictionary(out);
            }

            std::shared_ptr<DataType> get_value_type() const override {
                return indices_kernel_->get_value_type();
            }

            std::shared_ptr<DataType> dictionary_get_value_type() const {
                return dictionary_value_type_;
            }

            /// This can't be called more than once because DictionaryUnifier::get_result()
            /// can't be called more than once and produce the same output.
            turbo::Result<std::shared_ptr<Array>> dictionary() const {
                if (!first_dictionary_) {  // append was never called
                    return nullptr;
                }
                if (!dictionary_unifier_) {  // append was called only once
                    return first_dictionary_;
                }

                auto out_dict_type = first_dictionary_->type();
                std::shared_ptr<Array> out_dict;
                TURBO_RETURN_NOT_OK(dictionary_unifier_->get_result(&out_dict_type, &out_dict));
                return out_dict;
            }

        private:
            std::unique_ptr<HashKernel> indices_kernel_;
            std::shared_ptr<Array> first_dictionary_;
            std::shared_ptr<DataType> dictionary_value_type_;
            std::unique_ptr<DictionaryUnifier> dictionary_unifier_;
        };

// ----------------------------------------------------------------------
        template<typename HashKernel>
        turbo::Result<std::unique_ptr<KernelState>> HashInit(KernelContext *ctx,
                                                             const KernelInitArgs &args) {
            auto result = std::make_unique<HashKernel>(args.inputs[0].get_shared_ptr(), args.options,
                                                       ctx->memory_pool());
            TURBO_RETURN_NOT_OK(result->reset());
            // R build with openSUSE155 requires an explicit unique_ptr construction
            return std::unique_ptr<KernelState>(std::move(result));
        }

        template<typename Action>
        KernelInit GetHashInit(Type::type type_id) {
            // ARROW-8933: Generate only a single hash kernel per physical data
            // representation
            switch (type_id) {
                case Type::NA:
                    return HashInit<NullHashKernel<Action>>;
                case Type::BOOL:
                    return HashInit<RegularHashKernel<BooleanType, Action>>;
                case Type::INT8:
                case Type::UINT8:
                    return HashInit<RegularHashKernel<UInt8Type, Action>>;
                case Type::INT16:
                case Type::UINT16:
                    return HashInit<RegularHashKernel<UInt16Type, Action>>;
                case Type::INT32:
                case Type::UINT32:
                case Type::FP32:
                case Type::DATE32:
                case Type::TIME32:
                case Type::INTERVAL_MONTHS:
                    return HashInit<RegularHashKernel<UInt32Type, Action>>;
                case Type::INT64:
                case Type::UINT64:
                case Type::FP64:
                case Type::DATE64:
                case Type::TIME64:
                case Type::TIMESTAMP:
                case Type::DURATION:
                case Type::INTERVAL_DAY_TIME:
                    return HashInit<RegularHashKernel<UInt64Type, Action>>;
                case Type::BINARY:
                case Type::STRING:
                    return HashInit<RegularHashKernel<BinaryType, Action, std::string_view>>;
                case Type::LARGE_BINARY:
                case Type::LARGE_STRING:
                    return HashInit<RegularHashKernel<LargeBinaryType, Action, std::string_view>>;
                case Type::BINARY_VIEW:
                case Type::STRING_VIEW:
                    return HashInit<RegularHashKernel<BinaryViewType, Action, std::string_view>>;
                case Type::FIXED_SIZE_BINARY:
                case Type::DECIMAL128:
                case Type::DECIMAL256:
                    return HashInit<RegularHashKernel<FixedSizeBinaryType, Action, std::string_view>>;
                case Type::INTERVAL_MONTH_DAY_NANO:
                    return HashInit<RegularHashKernel<MonthDayNanoIntervalType, Action>>;
                default:
                    unreachable("non hashable type");
            }
        }

        using DictionaryEncodeState = OptionsWrapper<DictionaryEncodeOptions>;

        template<typename Action>
        turbo::Result<std::unique_ptr<KernelState>> DictionaryHashInit(KernelContext *ctx,
                                                                       const KernelInitArgs &args) {
            const auto &dict_type = turbo::checked_cast<const DictionaryType &>(*args.inputs[0].type);
            TURBO_MOVE_OR_RAISE(auto indices_hasher,
                                GetHashInit<Action>(dict_type.index_type()->id())(ctx, args));
            return std::make_unique<DictionaryHashKernel>(
                    turbo::checked_pointer_cast<HashKernel>(std::move(indices_hasher)),
                    dict_type.get_value_type());
        }

        turbo::Status HashExec(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
            auto hash_impl = turbo::checked_cast<HashKernel *>(ctx->state());
            TURBO_RETURN_NOT_OK(hash_impl->append(ctx, batch[0].array));
            TURBO_RETURN_NOT_OK(hash_impl->flush(out));
            return turbo::OkStatus();
        }

        turbo::Status UniqueFinalize(KernelContext *ctx, std::vector<Datum> *out) {
            auto hash_impl = turbo::checked_cast<HashKernel *>(ctx->state());
            std::shared_ptr<ArrayData> uniques;
            TURBO_RETURN_NOT_OK(hash_impl->GetDictionary(&uniques));
            *out = {Datum(uniques)};
            return turbo::OkStatus();
        }

        turbo::Status DictEncodeFinalize(KernelContext *ctx, std::vector<Datum> *out) {
            auto hash_impl = turbo::checked_cast<HashKernel *>(ctx->state());
            std::shared_ptr<ArrayData> uniques;
            TURBO_RETURN_NOT_OK(hash_impl->GetDictionary(&uniques));
            auto dict_type = dictionary(int32(), uniques->type);
            auto dict = make_array(uniques);
            for (size_t i = 0; i < out->size(); ++i) {
                (*out)[i] =
                        std::make_shared<DictionaryArray>(dict_type, (*out)[i].make_array(), dict);
            }
            return turbo::OkStatus();
        }

        std::shared_ptr<ArrayData> BoxValueCounts(const std::shared_ptr<ArrayData> &uniques,
                                                  const std::shared_ptr<ArrayData> &counts) {
            auto data_type =
                    STRUCT({field(kValuesFieldName, uniques->type), field(kCountsFieldName, int64())});
            ArrayVector children = {make_array(uniques), make_array(counts)};
            return std::make_shared<StructArray>(data_type, uniques->length, children)->data();
        }

        turbo::Status ValueCountsFinalize(KernelContext *ctx, std::vector<Datum> *out) {
            auto hash_impl = turbo::checked_cast<HashKernel *>(ctx->state());
            std::shared_ptr<ArrayData> uniques;

            TURBO_RETURN_NOT_OK(hash_impl->GetDictionary(&uniques));

            ExecResult result;
            TURBO_RETURN_NOT_OK(hash_impl->FlushFinal(&result));
            *out = {Datum(BoxValueCounts(uniques, result.array_data()))};
            return turbo::OkStatus();
        }

// Return the dictionary from the hash kernel or allocate an empty one.
// Required because on empty inputs, we don't ever see the input and
// hence have no dictionary.
        turbo::Result<std::shared_ptr<ArrayData>> EnsureHashDictionary(KernelContext *ctx,
                                                                       DictionaryHashKernel *hash) {
            TURBO_MOVE_OR_RAISE(auto dict, hash->dictionary());
            if (dict) {
                return dict->data();
            }
            TURBO_MOVE_OR_RAISE(auto null, MakeArrayOfNull(hash->dictionary_get_value_type(),
                    /*length=*/0, ctx->memory_pool()));
            return null->data();
        }

        turbo::Status UniqueFinalizeDictionary(KernelContext *ctx, std::vector<Datum> *out) {
            TURBO_RETURN_NOT_OK(UniqueFinalize(ctx, out));
            auto hash = turbo::checked_cast<DictionaryHashKernel *>(ctx->state());
            TURBO_MOVE_OR_RAISE((*out)[0].mutable_array()->dictionary,
                                EnsureHashDictionary(ctx, hash));
            return turbo::OkStatus();
        }

        turbo::Status ValueCountsFinalizeDictionary(KernelContext *ctx, std::vector<Datum> *out) {
            auto hash = turbo::checked_cast<DictionaryHashKernel *>(ctx->state());
            std::shared_ptr<ArrayData> uniques;
            ExecResult result;
            TURBO_RETURN_NOT_OK(hash->GetDictionary(&uniques));
            TURBO_RETURN_NOT_OK(hash->FlushFinal(&result));
            TURBO_MOVE_OR_RAISE(uniques->dictionary, EnsureHashDictionary(ctx, hash));
            *out = {Datum(BoxValueCounts(uniques, result.array_data()))};
            return turbo::OkStatus();
        }

        turbo::Result<TypeHolder> DictEncodeOutput(KernelContext *,
                                                   const std::vector<TypeHolder> &types) {
            return dictionary(int32(), types[0].get_shared_ptr());
        }

        turbo::Result<TypeHolder> ValueCountsOutput(KernelContext *,
                                                    const std::vector<TypeHolder> &types) {
            return STRUCT({field(kValuesFieldName, types[0].get_shared_ptr()),
                            field(kCountsFieldName, int64())});
        }

        template<typename Action>
        void AddHashKernels(VectorFunction *func, VectorKernel base, OutputType out_ty) {
            for (const auto &ty: primitive_types()) {
                base.init = GetHashInit<Action>(ty->id());
                base.signature = KernelSignature::create({ty}, out_ty);
                KCHECK_OK(func->add_kernel(base));
            }

            // Parametric types that we want matching to be dependent only on type id
            auto parametric_types = {Type::TIME32, Type::TIME64, Type::TIMESTAMP, Type::DURATION,
                                     Type::FIXED_SIZE_BINARY};
            for (const auto &type_id: parametric_types) {
                base.init = GetHashInit<Action>(type_id);
                base.signature = KernelSignature::create({type_id}, out_ty);
                KCHECK_OK(func->add_kernel(base));
            }

            for (auto t: {Type::DECIMAL128, Type::DECIMAL256}) {
                base.init = GetHashInit<Action>(t);
                base.signature = KernelSignature::create({t}, out_ty);
                KCHECK_OK(func->add_kernel(base));
            }

            for (const auto &ty: interval_types()) {
                base.init = GetHashInit<Action>(ty->id());
                base.signature = KernelSignature::create({ty}, out_ty);
                KCHECK_OK(func->add_kernel(base));
            }
        }

        const FunctionDoc unique_doc("Compute unique elements",
                                     ("Return an array with distinct values.\n"
                                      "Nulls are considered as a distinct value as well."),
                                     {"array"});

        const FunctionDoc value_counts_doc(
                "Compute counts of unique elements",
                ("For each distinct value, compute the number of times it occurs in the array.\n"
                 "The result is returned as an array of `struct<input type, int64>`.\n"
                 "Nulls in the input are counted and included in the output as well."),
                {"array"});

        const DictionaryEncodeOptions *GetDefaultDictionaryEncodeOptions() {
            static const auto kDefaultDictionaryEncodeOptions = DictionaryEncodeOptions::defaults();
            return &kDefaultDictionaryEncodeOptions;
        }

        const FunctionDoc dictionary_encode_doc(
                "Dictionary-encode array",
                ("Return a dictionary-encoded version of the input array.\n"
                 "This function does nothing if the input is already a dictionary array."),
                {"array"}, "DictionaryEncodeOptions");

// ----------------------------------------------------------------------
// This function does not use any hashing utilities
// but is kept in this file to be near dictionary_encode
// Dictionary decode implementation

        const FunctionDoc dictionary_decode_doc{
                "Decodes a DictionaryArray to an Array",
                ("Return a plain-encoded version of the array input\n"
                 "This function does nothing if the input is not a dictionary."),
                {"dictionary_array"}};

        class DictionaryDecodeMetaFunction : public MetaFunction {
        public:
            DictionaryDecodeMetaFunction()
                    : MetaFunction("dictionary_decode", Arity::Unary(), dictionary_decode_doc) {}

            turbo::Result<Datum> ExecuteImpl(const std::vector<Datum> &args,
                                             const FunctionOptions *options,
                                             ExecContext *ctx) const override {
                if (args[0].type() == nullptr || args[0].type()->id() != Type::DICTIONARY) {
                    return args[0];
                }

                if (args[0].is_array() || args[0].is_chunked_array()) {
                    DictionaryType *dict_type = turbo::checked_cast<DictionaryType *>(args[0].type().get());
                    CastOptions cast_options = CastOptions::Safe(dict_type->get_value_type());
                    return call_function("cast", args, &cast_options, ctx);
                } else {
                    return turbo::failed_precondition_error("Expected an Array or a Chunked Array");
                }
            }
        };
    }  // namespace

    void register_vector_hash(FunctionRegistry *registry) {
        VectorKernel base;
        base.exec = HashExec;

        // ----------------------------------------------------------------------
        // unique

        base.finalize = UniqueFinalize;
        base.output_chunked = false;
        auto unique = std::make_shared<VectorFunction>("unique", Arity::Unary(), unique_doc);
        AddHashKernels<UniqueAction>(unique.get(), base, FirstType);

        // Dictionary unique
        base.init = DictionaryHashInit<UniqueAction>;
        base.finalize = UniqueFinalizeDictionary;
        base.signature = KernelSignature::create({Type::DICTIONARY}, FirstType);
        KCHECK_OK(unique->add_kernel(base));

        KCHECK_OK(registry->add_function(std::move(unique)));

        // ----------------------------------------------------------------------
        // value_counts

        base.finalize = ValueCountsFinalize;
        auto value_counts =
                std::make_shared<VectorFunction>("value_counts", Arity::Unary(), value_counts_doc);
        AddHashKernels<ValueCountsAction>(value_counts.get(), base, ValueCountsOutput);

        // Dictionary value counts
        base.init = DictionaryHashInit<ValueCountsAction>;
        base.finalize = ValueCountsFinalizeDictionary;
        base.signature = KernelSignature::create({Type::DICTIONARY}, ValueCountsOutput);
        KCHECK_OK(value_counts->add_kernel(base));

        KCHECK_OK(registry->add_function(std::move(value_counts)));

        // ----------------------------------------------------------------------
        // dictionary_encode

        base.finalize = DictEncodeFinalize;
        // Unique and ValueCounts output unchunked arrays
        base.output_chunked = true;

        auto dict_encode = std::make_shared<VectorFunction>(
                "dictionary_encode", Arity::Unary(), dictionary_encode_doc,
                GetDefaultDictionaryEncodeOptions());
        AddHashKernels<DictEncodeAction>(dict_encode.get(), base, DictEncodeOutput);

        auto no_op = [](KernelContext *, const ExecSpan &span, ExecResult *out) {
            out->value = span[0].array.to_array_data();
            return turbo::OkStatus();
        };
        KCHECK_OK(dict_encode->add_kernel({Type::DICTIONARY}, OutputType(FirstType), no_op));

        KCHECK_OK(registry->add_function(std::move(dict_encode)));
    }

    void register_dictionary_decode(FunctionRegistry *registry) {
        KCHECK_OK(registry->add_function(std::make_shared<DictionaryDecodeMetaFunction>()));
    }

}  // namespace nebula::compute::internal
